import cv2
import os
from threading import Thread, Lock
import queue
from ultralytics import YOLO

# 防止 OpenMP 库冲突（Windows）
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# 加载训练好的 YOLO 模型
model_path = r"E:\qwj_work\ultralytics\ultralytics\runs\result\LD.2\weights\best.pt"
model = YOLO(model_path)

# RTSP 视频流地址
rtsp_url = "rtsp://admin:EJKUIB@192.168.31.63:554/h264/ch1/sub/av_stream"

# 设置抽帧参数
skip_frames = 6  # 每 6 帧处理一次

# 设置推理输入大小
img_size = 640

# 创建队列用于帧传递
raw_frame_queue = queue.Queue(maxsize=10)  # 原始帧队列
result_frame_queue = queue.Queue(maxsize=10)  # 检测结果队列

# 视频读取线程
def video_reader():
    cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG)
    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print("❌ 无法接收帧，可能视频流已断开...")
            break

        frame_count += 1
        if frame_count % skip_frames == 0:
            if raw_frame_queue.full():
                raw_frame_queue.get()  # 丢弃旧帧
            raw_frame_queue.put(frame)

    cap.release()
    raw_frame_queue.put(None)  # 发送结束信号
    result_frame_queue.put(None)

# 推理线程
def inference_worker():
    while True:
        if raw_frame_queue.empty():
            continue
        frame = raw_frame_queue.get()
        if frame is None:
            result_frame_queue.put(None)
            break

        # 推理
        try:
            results = model(frame, imgsz=img_size, conf=0.25, half=True, stream=True)
            for result in results:
                annotated_frame = result.plot()
                annotated_frame = cv2.resize(annotated_frame, (768, 432))
                if not result_frame_queue.full():
                    result_frame_queue.put(annotated_frame)
        except Exception as e:
            print("⚠️ 推理异常:", e)

# 启动线程
Thread(target=video_reader, daemon=True).start()
Thread(target=inference_worker, daemon=True).start()

# 创建显示窗口
cv2.namedWindow("YOLOv11 RTSP Detection", cv2.WINDOW_NORMAL)
cv2.resizeWindow("YOLOv11 RTSP Detection", 768, 432)

# 主线程负责显示
while True:
    if not result_frame_queue.empty():
        annotated_frame = result_frame_queue.get()
        if annotated_frame is None:
            print("⚠️ 推理线程已结束")
            break

        cv2.imshow("YOLOv8 RTSP Detection", annotated_frame)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
    else:
        # 可选：防止 CPU 占用过高
        cv2.waitKey(1)

# 释放资源
cv2.destroyAllWindows()